import re
import logging

from agent.common.system import sysCommand
from agent.feature.base import Feature
from agent.common.netinfo import NetInfo
from agent.common.config import Config

logger = logging.getLogger('common')

class SmpAffinity(Feature):
    __domain__  = "net"
    __feature__ = "smp_affinity"
    def __init__(self) -> None:
        super().__init__(self.__domain__, self.__feature__)
        self.net_info = NetInfo()
        assert self.net_info.numa_node_core_range.__len__() != 0
    
    def set_value(self, value):
        if value == "off":
            sysCommand("systemctl start irqbalance.service",  log = True)
            sysCommand("systemctl enable irqbalance.service",  log = True)
        else:
            sysCommand("systemctl stop irqbalance.service",  log = True)
            sysCommand("systemctl disable irqbalance.service", log = True)
        
        if type(value) == str and re.match(r"\d+-\d+",value):
            # setting speccpu affinity by set_irq.sh
            sysCommand("sh set_irq.sh {eth_name} speccpu '{cpu_range}'".format(
                eth_name = self.net_info.dev_name,
                cpu_range = value
            ), cwd = Config.SCRIPTS_PATH, log = True)
        
        if type(value) == str and value in ['numa0', 'numa1', 'balance']:
            # setting affinity by set_irq.sh
            sysCommand("sh set_irq.sh {eth_name} {value}".format(
                eth_name = self.net_info.dev_name,
                value = value
            ), cwd = Config.SCRIPTS_PATH, log = True)
        
        elif type(value) == str and value in ['dedicated', 'disperse', 'different', 'off', 'eighth']:
            # setting affinity by netinfo
            affinity_bingding_schema = self.net_info.get_affinity_binding_schema(value)
            for interrupt_index, core_num in affinity_bingding_schema.items():
                sysCommand("echo {} > /proc/irq/{}/smp_affinity_list".format(
                    core_num, interrupt_index))
        
        elif type(value) == dict:
            # setting backup value
            for interrupt_index, core_num in value.items():
                sysCommand("echo {} > /proc/irq/{}/smp_affinity_list".format(
                    core_num, interrupt_index))
            
    def get_value(self):
        affinity_bingding_schema = {}
        for interrupt_index in self.net_info.interrupts:
            affinity_bingding_schema[interrupt_index] = sysCommand(
                "cat /proc/irq/{}/smp_affinity_list".format(interrupt_index))
        return affinity_bingding_schema
        
if __name__ == "__main__":
    handler = SmpAffinity()
    handler.backup()
    handler.set_value("eighth")
    handler.set_value("0-3")
    handler.set_value("numa0")
    handler.set_value("balance")
    handler.rollback()

